# -*- coding: utf-8 -*-
# 功能介绍:
    # load_image(image_path)
    # save_image_and_label(image, labels, image_name, label_name):
    # augment_image_and_label(image, labels):
    # augment_image_and_label1(image, labels, rotation_range=30, shift_range=0.1, zoom_range=0.2, flip_prob=0.5):
    

import os
import cv2
import numpy as np
import random
from PIL import Image

# 路径设置
image_dir = "D:/上过的课/2024-2025-1-大四上/毕设/datasets/washed1/images"  # 图像目录
label_dir = "D:/上过的课/2024-2025-1-大四上/毕设/datasets/washed1/labels"  # 标注目录
output_image_dir = "D:/上过的课/2024-2025-1-大四上/毕设/datasets/augmented/images"  # 增强后图像保存目录
output_label_dir = "D:/上过的课/2024-2025-1-大四上/毕设/datasets/augmented/labels"  # 增强后标注保存目录

# 创建输出目录
os.makedirs(output_image_dir, exist_ok=True)
os.makedirs(output_label_dir, exist_ok=True)

# 数据增强参数
num_augmentations = 5  # 每张类型1的图像增强的次数
rotation_range = 15  # 随机旋转角度范围
shift_range = 0.1  # 随机平移范围
zoom_range = 0.2  # 随机缩放范围
flip_prob = 0.5  # 随机翻转概率

# 使用 cv2.imdecode() 读取图像，避免中文路径加载失败
def load_image(image_path):
    with open(image_path, "rb") as f:
        image_data = np.frombuffer(f.read(), dtype=np.uint8)
    image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    return image

def save_image_and_label(image, labels, image_name, label_name):
    """保存增强后的图像和标注"""
    # 检查图像是否为 None
    if image is None:
        print(f"Augmented image is None, skipping save for {image_name}")
        return
    
    # 保存图像
    output_image_path = os.path.join(output_image_dir, image_name)
    # print(f"Saving image to: {output_image_path}")
    success, encoded_image = cv2.imencode(".jpg", image)  # 编码为 .jpg 格式，imencode避免中文路径加载失败
    if not success:
        print(f"Failed to encode image: {output_image_path}")
        return
    
    # 将二进制数据写入文件
    with open(output_image_path, "wb") as f:
        f.write(encoded_image.tobytes())
    
    # 保存标注
    output_label_path = os.path.join(output_label_dir, label_name)
    with open(output_label_path, "w", encoding="utf-8") as f:
        for label in labels:
            f.write("".join(label) + "\n")

def augment_image_and_label(image, labels, rotation_range=30, shift_range=0.1, zoom_range=0.2, flip_prob=0.5):
    """对图像和标注进行数据增强，并同步更新标注框坐标"""
    h, w = image.shape[:2]

    # 保存原始尺寸用于最后的尺寸恢复
    original_h, original_w = h, w

    # 先执行缩放操作（保持宽高比）
    scale = random.uniform(1 - zoom_range, 1 + zoom_range)
    new_w = int(w * scale)
    new_h = int(h * scale)
    image = cv2.resize(image, (new_w, new_h))
    h, w = new_h, new_w  # 更新当前尺寸

    # 随机平移（在缩放后的尺寸上进行）
    shift_x = random.uniform(-shift_range, shift_range) * w
    shift_y = random.uniform(-shift_range, shift_range) * h
    M = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
    image = cv2.warpAffine(image, M, (w, h))

    # 调整平移后的标注坐标
    valid_labels = []
    for label in labels:
        class_id, xc, yc, bw, bh = map(float, label)
        
        # 转换为像素坐标并应用平移
        xc = xc * w + shift_x
        yc = yc * h + shift_y

        # 转换回归一化坐标
        valid_labels.append([class_id, xc/w, yc/h, bw, bh])
    # 剔除中心点跑到图外的标签
    labels = []
    for label in valid_labels:
        class_id, xc, yc, bw, bh = label
        if (xc > 0 and yc > 0 and xc < 1 and yc < 1):
            labels.append(label)

    # 随机旋转（在平移后的坐标上进行）
    angle = random.uniform(-rotation_range, rotation_range)
    M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1)
    image = cv2.warpAffine(image, M, (w, h))

    # 调整旋转后的标注坐标
    valid_labels = []
    for label in labels:
        class_id, xc, yc, bw, bh = map(float, label)
        xc = xc * w
        yc = yc * h
        point = np.array([[xc, yc, 1]])  # 齐次坐标 [x, y, 1]
        M_homogeneous = np.vstack([M, [0, 0, 1]])  # 扩展为 3x3
        # 计算变换后的点
        transformed_point_homogeneous = M_homogeneous @ point.T  # 3x3 @ 3x1 → 3x1
        transformed_point = transformed_point_homogeneous[:2].T / transformed_point_homogeneous[2]  # 归一化
        # 提取结果
        transformed_point = transformed_point[0]  # [x', y']
        xc=transformed_point[0]
        yc=transformed_point[1]
        valid_labels.append([class_id, xc/w, yc/h, bw, bh])
    # 剔除中心点跑到图外的标签
    labels = []
    for label in valid_labels:
        class_id, xc, yc, bw, bh = label
        if (xc > 0 and yc > 0 and xc < 1 and yc < 1):
            labels.append(label)

    # 随机水平翻转
    if random.random() < flip_prob:
        image = cv2.flip(image, 1)
        labels = [
            [class_id, 1.0 - xc, yc, bw, bh]
            for class_id, xc, yc, bw, bh in labels
        ]

    # 将labels转换为字符串格式
    str_labels = []
    for label in labels:
        # 将类别保存为整数，其他值保存为浮点数
        class_id = int(label[0])  # 类别转换为整数
        coordinates = list(map(str, label[1:]))  # 其他值转换为字符串
        str_label = f"{class_id} {' '.join(coordinates)}"  # 组合成字符串
        str_labels.append(str_label)

    return image, str_labels

def count():
    #统计和预测数据增强影响范围
    a0 = 0;
    a1 = 0;
    a2 = 0
    b0 = 0;
    b1 = 0;
    b2 = 0
    c0 = 0;
    c1 = 0;
    c2 = 0
    # 遍历标注文件
    for label_file in os.listdir(label_dir):
        label_path = os.path.join(label_dir, label_file)
        with open(label_path, "r") as f:
            labels = [line.strip().split() for line in f.readlines()]
        
        # 统计类型0、类型1和类型2的标签数量
        count0 = sum(1 for label in labels if int(label[0]) == 0)
        count1 = sum(1 for label in labels if int(label[0]) == 1)
        count2 = sum(1 for label in labels if int(label[0]) == 2)
        # a 只选择类型1的标签数大于类型0和类型2标签数之和的图像
        if (count1 > (count0 + count2)):
            a0+=count0
            a1+=count1
            a2+=count2
           # print(count0, count1, count2, label_file,sep=' ')
        # b 只选择类型1标签数>类型0>类型2的图像
        if(count1>count0 and count0>count2):
            b0+=count0
            b1+=count1
            b2+=count2
        # c 只选择 a且b的图像
        if(count1>(count0 + count2) and count0>count2):
            c0+=count0
            c1+=count1
            c2+=count2
    print("a型：",a0,a1,a2)
    print("b型：",b0,b1,b2)
    print("c型：",c0,c1,c2)

def main():
    # 遍历标注文件
    for label_file in os.listdir(label_dir):
        label_path = os.path.join(label_dir, label_file)
        with open(label_path, "r") as f:
            labels = [line.strip().split() for line in f.readlines()]
        
        # 统计类型0、类型1和类型2的标签数量
        count_0 = sum(1 for label in labels if int(label[0]) == 0)
        count_1 = sum(1 for label in labels if int(label[0]) == 1)
        count_2 = sum(1 for label in labels if int(label[0]) == 2)
        # 只选择类型1的标签数大于类型0和类型2标签数之和的图像
        if not (count_1 > (count_0 + count_2)):
            continue
        

        # 加载对应的图像
        image_file = label_file.replace(".txt", ".jpg")
        image_path = os.path.join(image_dir, image_file)
        if not os.path.exists(image_path):
            print(f"Image {image_file} not found, skipping.")
            continue

        image = load_image(image_path)

        # 对类型1的样本进行数据增强
        for i in range(num_augmentations):
            augmented_image, augmented_labels = augment_image_and_label(image.copy(), labels.copy())
            # 保存增强后的图像和标注
            output_image_name = f"{os.path.splitext(image_file)[0]}_aug_{i}.jpg"
            output_label_name = f"{os.path.splitext(label_file)[0]}_aug_{i}.txt"
            save_image_and_label(augmented_image, augmented_labels, output_image_name, output_label_name)


if __name__ == "__main__":
    main()